/**
 *  Copyright 2006 Juan Manuel Caicedo 
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License"); 
 *  you may not use this file except in compliance with the License. 
 *  You may obtain a copy of the License at 
 *  
 *  http://www.apache.org/licenses/LICENSE-2.0 
 *  Unless required by applicable law or agreed to in writing, software 
 *  distributed under the License is distributed on an "AS IS" BASIS, 
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 *  See the License for the specific language governing permissions and 
 *  limitations under the License.
 * 
 */
package gate.creole.opennlp;

import gate.Annotation;
import gate.AnnotationSet;
import gate.Factory;
import gate.FeatureMap;
import gate.Resource;
import gate.creole.AbstractLanguageAnalyser;
import gate.creole.ExecutionException;
import gate.creole.ResourceInstantiationException;
import gate.util.GateRuntimeException;
import gate.util.OffsetComparator;
import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.ListIterator;
import opennlp.tools.ngram.Dictionary;
import opennlp.tools.postag.POSDictionary;
import opennlp.tools.postag.TagDictionary;

/**
 * 
 * 
 * @author Juan Manuel Caicedo
 * 
 */
public class POSTagger extends AbstractLanguageAnalyser {
  public static final String TAG_DOCUMENT_PARAMETER_NAME = "document";

  public static final String TAG_INPUT_AS_PARAMETER_NAME = "inputASName";

  public static final String TAG_LEXICON_URL_PARAMETER_NAME = "lexiconURL";

  public static final String TAG_RULES_URL_PARAMETER_NAME = "rulesURL";

  public static final String TAG_ENCODING_PARAMETER_NAME = "encoding";

  public static final String BASE_TOKEN_ANNOTATION_TYPE_PARAMETER_NAME = "baseTokenAnnotationType";

  public static final String OUTPUT_ANNOTATION_TYPE_PARAMETER_NAME = "outputAnnotationType";

  public static final String BASE_SENTENCE_ANNOTATION_TYPE_PARAMETER_NAME = "baseSentenceAnnotationType";

  public static final String TAG_OUTPUT_AS_PARAMETER_NAME = "outputASName";

  public POSTagger() {
  }

  public Resource init() throws ResourceInstantiationException {
    if(modelURL == null) { throw new ResourceInstantiationException(
            "No URL provided for the model!"); }
    if(language == null)
      throw new ResourceInstantiationException("No Language provided");
    try {
      TagDictionary tagDict = null;
      if(tagDictionaryURL != null) {
        tagDict = new POSDictionary(Util.absoluteFilePath(tagDictionaryURL));
      }
      String modelPath = Util.absoluteFilePath(modelURL);
      if(language.equals("spanish")) {
        tagger = (tagDict != null) ? new opennlp.tools.lang.spanish.PosTagger(
                modelPath, tagDict) : new opennlp.tools.lang.spanish.PosTagger(
                modelPath);
      } else if(language.equals("english")) {
        if(dictionaryURL == null && tagDictionaryURL == null)
          throw new ResourceInstantiationException(
                  "At leat a Dictionary must be provided");
        if(dictionaryURL != null) {
          Dictionary dict = new Dictionary(Util.absoluteFilePath(dictionaryURL));
          tagger = (tagDict != null)
                  ? new opennlp.tools.lang.english.PosTagger(modelPath, dict,
                          tagDict)
                  : new opennlp.tools.lang.english.PosTagger(modelPath, dict);
        } else {
          tagger = new opennlp.tools.lang.english.PosTagger(modelPath, tagDict);
        }
      }
    } catch(Exception e) {
      throw new ResourceInstantiationException(e);
    }
    return this;
  }

  public void execute() throws ExecutionException {
    try {
      // check the parameters
      if(document == null)
        throw new GateRuntimeException("No document to process!");
      if(inputASName != null && inputASName.equals("")) inputASName = null;
      AnnotationSet inputAS = (inputASName == null)
              ? document.getAnnotations()
              : document.getAnnotations(inputASName);
      /* Addition by Niraj */
      if(baseTokenAnnotationType == null
              || baseTokenAnnotationType.trim().length() == 0) { throw new GateRuntimeException(
              "No base Token Annotation Type provided!"); }
      if(outputASName != null && outputASName.equals("")) outputASName = null;
      AnnotationSet outputAS = (outputASName == null) ? document
              .getAnnotations() : document.getAnnotations(outputASName);
      if(baseSentenceAnnotationType == null
              || baseSentenceAnnotationType.trim().length() == 0) { throw new GateRuntimeException(
              "No base Sentence Annotation Type provided!"); }
      if(outputAnnotationType == null
              || outputAnnotationType.trim().length() == 0) { throw new GateRuntimeException(
              "No AnnotationType provided to store the new feature!"); }
      /* End of addition */
      AnnotationSet sentencesAS = inputAS.get(baseSentenceAnnotationType);
      AnnotationSet tokensAS = inputAS.get(baseTokenAnnotationType);
      if(sentencesAS != null && sentencesAS.size() > 0 && tokensAS != null
              && tokensAS.size() > 0) {
        long startTime = System.currentTimeMillis();
        fireStatusChanged("POS tagging " + document.getName());
        fireProgressChanged(0);
        // prepare the input for HepTag
        List sentenceForTagger = new ArrayList();
        List sentencesForTagger = new ArrayList(1);
        sentencesForTagger.add(sentenceForTagger);
        // define a comparator for annotations by start offset
        Comparator offsetComparator = new OffsetComparator();
        // read all the tokens and all the sentences
        List sentencesList = new ArrayList(sentencesAS);
        Collections.sort(sentencesList, offsetComparator);
        List tokensList = new ArrayList(tokensAS);
        Collections.sort(tokensList, offsetComparator);
        Iterator sentencesIter = sentencesList.iterator();
        ListIterator tokensIter = tokensList.listIterator();
        List tokensInCurrentSentence = new ArrayList();
        Annotation currentToken = (Annotation)tokensIter.next();
        int sentIndex = 0;
        int sentCnt = sentencesAS.size();
        while(sentencesIter.hasNext()) {
          Annotation currentSentence = (Annotation)sentencesIter.next();
          tokensInCurrentSentence.clear();
          sentenceForTagger.clear();
          while(currentToken != null
                  && currentToken.getEndNode().getOffset().compareTo(
                          currentSentence.getEndNode().getOffset()) <= 0) {
            tokensInCurrentSentence.add(currentToken);
            sentenceForTagger.add(currentToken.getFeatures().get(
                    TOKEN_STRING_FEATURE_NAME));
            currentToken = (Annotation)(tokensIter.hasNext() ? tokensIter
                    .next() : null);
          }
          // run the POS tagger
          List taggerResults = tagger.tag(sentenceForTagger);
          if(taggerResults.size() != tokensInCurrentSentence.size())
            throw new GateRuntimeException(
                    "POS Tagger malfunction: the output size ("
                            + taggerResults.size()
                            + ") is different from the input size ("
                            + tokensInCurrentSentence.size() + ")!");
          Iterator tokIter = tokensInCurrentSentence.iterator();
          Iterator tagIter = taggerResults.iterator();
          while(tagIter.hasNext()) {
            String tag = (String)tagIter.next();
            Annotation annot = (Annotation)tokIter.next();
            addFeatures(annot, TOKEN_CATEGORY_FEATURE_NAME, tag);
          }
          fireProgressChanged(sentIndex++ * 100 / sentCnt);
        }// while(sentencesIter.hasNext())
        if(currentToken != null) {
          // we have remaining tokens after the last sentence
          tokensInCurrentSentence.clear();
          sentenceForTagger.clear();
          while(currentToken != null) {
            tokensInCurrentSentence.add(currentToken);
            sentenceForTagger.add(currentToken.getFeatures().get(
                    TOKEN_STRING_FEATURE_NAME));
            currentToken = (Annotation)(tokensIter.hasNext() ? tokensIter
                    .next() : null);
          }
          List taggerResults = tagger.tag(sentenceForTagger);
          if(taggerResults.size() != tokensInCurrentSentence.size())
            throw new GateRuntimeException(
                    "POS Tagger malfunction: the output size ("
                            + taggerResults.size()
                            + ") is different from the input size ("
                            + tokensInCurrentSentence.size() + ")!");
          Iterator tokIter = tokensInCurrentSentence.iterator();
          Iterator tagIter = taggerResults.iterator();
          while(tagIter.hasNext()) {
            String tag = (String)tagIter.next();
            Annotation annot = (Annotation)tokIter.next();
            addFeatures(annot, TOKEN_CATEGORY_FEATURE_NAME, tag);
          }
        }// if(currentToken != null)
        fireProcessFinished();
        fireStatusChanged(document.getName()
                + " tagged in "
                + NumberFormat
                        .getInstance()
                        .format(
                                (double)(System.currentTimeMillis() - startTime) / 1000)
                + " seconds!");
      } else {
        throw new GateRuntimeException("No sentences or tokens to process!\n"
                + "Please run a sentence splitter " + "and tokeniser first!");
      }
    } catch(Exception e) {
      throw new ExecutionException(e);
    }
  }

  protected void addFeatures(Annotation annot, String featureName,
          String featureValue) throws GateRuntimeException {
    String tempIASN = inputASName == null ? "" : inputASName;
    String tempOASN = outputASName == null ? "" : outputASName;
    if(outputAnnotationType.equals(baseTokenAnnotationType)
            && tempIASN.equals(tempOASN)) {
      annot.getFeatures().put(featureName, featureValue);
      return;
    } else {
      int start = annot.getStartNode().getOffset().intValue();
      int end = annot.getEndNode().getOffset().intValue();
      // get the annotations of type outputAnnotationType
      AnnotationSet outputAS = (outputASName == null) ? document
              .getAnnotations() : document.getAnnotations(outputASName);
      AnnotationSet annotations = outputAS.get(outputAnnotationType);
      if(annotations == null || annotations.size() == 0) {
        // add new annotation
        FeatureMap features = Factory.newFeatureMap();
        features.put(featureName, featureValue);
        try {
          outputAS.add(new Long(start), new Long(end), outputAnnotationType,
                  features);
        } catch(Exception e) {
          throw new GateRuntimeException("Invalid Offsets");
        }
      } else {
        // search for the annotation if there is one with the same start
        // and end offsets
        ArrayList tempList = new ArrayList(annotations.get());
        boolean found = false;
        for(int i = 0; i < tempList.size(); i++) {
          Annotation annotation = (Annotation)tempList.get(i);
          if(annotation.getStartNode().getOffset().intValue() == start
                  && annotation.getEndNode().getOffset().intValue() == end) {
            // this is the one
            annotation.getFeatures().put(featureName, featureValue);
            found = true;
            break;
          }
        }
        if(!found) {
          // add new annotation
          FeatureMap features = Factory.newFeatureMap();
          features.put(featureName, featureValue);
          try {
            outputAS.add(new Long(start), new Long(end), outputAnnotationType,
                    features);
          } catch(Exception e) {
            throw new GateRuntimeException("Invalid Offsets");
          }
        }
      }
    }
  }

  public String getLanguage() {
    return language;
  }

  public void setLanguage(String language) {
    this.language = language;
  }

  public java.net.URL getModelURL() {
    return modelURL;
  }

  public void setModelURL(java.net.URL modelURL) {
    this.modelURL = modelURL;
  }

  public java.net.URL getDictionaryURL() {
    return dictionaryURL;
  }

  public void setDictionaryURL(java.net.URL dictionaryURL) {
    this.dictionaryURL = dictionaryURL;
  }

  public java.net.URL getTagDictionaryURL() {
    return tagDictionaryURL;
  }

  public void setTagDictionaryURL(java.net.URL tagDictionaryURL) {
    this.tagDictionaryURL = tagDictionaryURL;
  }

  public void setInputASName(String newInputASName) {
    inputASName = newInputASName;
  }

  public String getInputASName() {
    return inputASName;
  }

  public String getBaseTokenAnnotationType() {
    return this.baseTokenAnnotationType;
  }

  public String getBaseSentenceAnnotationType() {
    return this.baseSentenceAnnotationType;
  }

  public String getOutputAnnotationType() {
    return this.outputAnnotationType;
  }

  public void setBaseTokenAnnotationType(String baseTokenAnnotationType) {
    this.baseTokenAnnotationType = baseTokenAnnotationType;
  }

  public void setBaseSentenceAnnotationType(String baseSentenceAnnotationtype) {
    this.baseSentenceAnnotationType = baseSentenceAnnotationtype;
  }

  public void setOutputAnnotationType(String outputAnnotationType) {
    this.outputAnnotationType = outputAnnotationType;
  }

  public String getOutputASName() {
    return this.outputASName;
  }

  public void setOutputASName(String outputASName) {
    this.outputASName = outputASName;
  }

  protected opennlp.tools.postag.POSTagger tagger;

  private String language;

  private java.net.URL modelURL;

  private java.net.URL dictionaryURL;

  private java.net.URL tagDictionaryURL;

  private String inputASName;

  /* Addition by Niraj */
  private String baseTokenAnnotationType;

  private String baseSentenceAnnotationType;

  private String outputAnnotationType;

  private String outputASName;
  /* End of Addition */
}
